package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/segmentio/kafka-go"
)

func main() {
	// 配置参数（需与消费者匹配）
	brokers := []string{"localhost:9091"}
	topic := "first_topic"

	// 创建Kafka写入器
	w := &kafka.Writer{
		Addr:                   kafka.TCP(brokers...),
		Topic:                  topic,
		BatchSize:              1,                    // 关键配置：每条消息立即发送
		BatchTimeout:           1 * time.Millisecond, // 最大批处理等待时间
		Balancer:               &kafka.LeastBytes{},  // 负载均衡策略
		RequiredAcks:           kafka.RequireAll,     // 等待所有副本确认
		AllowAutoTopicCreation: false,                // 禁用自动创建主题
		WriteTimeout:           10 * time.Second,     // 写入超时
		ReadTimeout:            10 * time.Second,     // 读取超时
	}

	// 启动消息生成协程
	ctx, cancel := context.WithCancel(context.Background())
	go generateMessages(ctx, w)

	// 优雅关闭
	waitForShutdown(w, cancel)
}

// 生成测试消息
func generateMessages(ctx context.Context, w *kafka.Writer) {
	counter := 0
	ticker := time.NewTicker(100 * time.Millisecond) // 每秒生成一条消息
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			counter++
			message := &kafka.Message{
				Key:   []byte(time.Now().Format("15:04:05")),
				Value: []byte(fmt.Sprintf("测试消息 #%d - %s", counter, time.Now().String())),
			}

			// 发送消息
			err := w.WriteMessages(ctx, *message)
			if err != nil {
				log.Printf("发送消息失败: %v", err)
				continue
			}
			log.Printf("已发送消息 #%d", counter)
		}
	}
}

// 处理程序退出逻辑
func waitForShutdown(w *kafka.Writer, cancel context.CancelFunc) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	log.Println("正在停止生产者...")
	cancel() // 停止消息生成

	// 确保所有消息都被发送
	_, timeout := context.WithTimeout(context.Background(), 10*time.Second)
	defer timeout()

	if err := w.Close(); err != nil {
		log.Fatalf("关闭写入器失败: %v", err)
	}
	log.Println("生产者已安全退出")
}
